import openai
import faiss
import numpy as np
import sqlite3
import json
import os

# 设置OpenAI API密钥
os.environ[
    "OPENAI_API_KEY"
] = "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"
client = openai.OpenAI()

# 文本
texts = [
    "sex on the beach的配方是2 oz. vodka	+ 1 oz peach liqueur + 1 oz. Spicy agave syrup and jalapeño	+ 1.5 orange juice + 1.5 Cranberry + Ice",
    "Christmas Martini 圣诞节马提尼 的配方是 1.5 oz vodka	1 .5 oz Baileys	2 oz. Coconut cream	1 oz. Amaretto	Ice",
    "Blue Ja Guayana 蓝色瓜亚那 的配方是 2 oz pineapple juice	 0.75 oz coconut cream	0.75 oz Blue curaçao 2 oz white rum	Ice",
    "Christmas Martini 是一款以圣诞节日氛围为主题的创意马天尼鸡尾酒，通常以红色、绿色、金色等圣诞配色呈现，并融入冬季风味（如肉桂、香草、蔓越莓等）。",
    "Christmas Martini 做法：摇壶中加入所有材料与冰块，摇匀后过滤倒入马天尼杯。用肉桂棒和迷迭香做装饰，撒少量肉豆蔻粉增加香气。",
    "Christmas Martini 视觉：鲜艳的节日色彩（红/绿/金），常用糖边、 edible glitter（可食用亮粉）增强氛围。",
    "Christmas Martini 味觉：甜味主导，带有香料（肉桂、肉豆蔻）、浆果或巧克力风味，酒精度中等（约15%-20%）。",
]


def get_openai_embeddings(texts, model="text-embedding-3-small"):
    """使用OpenAI API获取文本嵌入向量"""
    response = client.embeddings.create(input=texts, model=model)
    return np.array([data.embedding for data in response.data])


# 生成嵌入向量
embeddings = get_openai_embeddings(texts)
dimension = embeddings.shape[1]  # 获取向量维度

# 创建FAISS索引
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)

# 保存FAISS索引到文件
faiss.write_index(index, "openai_rag_index.faiss")

# 创建SQLite数据库存储元数据
conn = sqlite3.connect("openai_rag_database.db")
cursor = conn.cursor()

# 创建表
cursor.execute(
    """
CREATE TABLE IF NOT EXISTS documents (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    text TEXT NOT NULL,
    embedding BLOB NOT NULL,
    metadata TEXT,
    model TEXT NOT NULL
)
"""
)

# 插入数据
for i, (text, embedding) in enumerate(zip(texts, embeddings)):
    # 将numpy数组转换为字节流
    embedding_bytes = embedding.tobytes()
    # 插入记录
    cursor.execute(
        """
    INSERT INTO documents (text, embedding, metadata, model)
    VALUES (?, ?, ?, ?)
    """,
        (text, embedding_bytes, json.dumps({"source": "example"}), "text-embedding-3-small"),
    )

conn.commit()
conn.close()

print("使用OpenAI模型生成的RAG向量已成功存储到数据库")
